/*
 * BeastMC3.java
 *
 * Copyright (C) 2002-2006 Alexei Drummond and Andrew Rambaut
 *
 * This file is part of BEAST.
 * See the NOTICE file distributed with this work for additional
 * information regarding copyright ownership and licensing.
 *
 * BEAST is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 *  BEAST is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with BEAST; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301  USA
 */

package dr.app.beastdev;

import dr.app.util.Arguments;
import dr.app.util.Utils;
import dr.inference.mcmc.MCMC;
import dr.inference.mcmcmc.MCMCMC;
import dr.inference.mcmcmc.MCMCMCOptions;
import dr.xml.ObjectStore;
import dr.xml.XMLParser;
import dr.util.MessageLogHandler;
import dr.util.Version;
import org.virion.jam.console.ConsoleApplication;

import java.io.*;
import java.util.Iterator;
import java.util.logging.*;

public class BeastMC3 {

    public final static int HOT_CHAIN_COUNT = 2;

	private final static Version version = new BeastVersion();

	static class BeastConsoleApp extends ConsoleApplication {
		XMLParser parser = null;

		public BeastConsoleApp(String nameString, String aboutString, javax.swing.Icon icon) throws IOException {
			super(nameString, aboutString, icon, false);
		}

		public void doStop() {
			Iterator iter = parser.getThreads();
			while (iter.hasNext()) {
				Thread thread = (Thread)iter.next();
				// @todo should never use this method...
                thread.stop();
                // @todo who cares? - it is bad in complex multithreading applications but this is a really simple case.
                // At the moment it works and when it doesn't we can implement a polling alternative.
            }
		}
	};

	public BeastMC3(double[] chainTemperatures, int swapChainsEvery, File inputFile, BeastConsoleApp consoleApp, boolean verbose) {

		if (inputFile == null) {
			System.err.println();
			System.err.println("Error: no input file specified");
			return;
		}

        int chainCount = chainTemperatures.length;
        MCMC[] chains = new MCMC[chainCount];
		MCMCMCOptions options = new MCMCMCOptions();
		options.setChainTemperatures(chainTemperatures);
		options.setSwapChainsEvery(swapChainsEvery);


		// Add a handler to handle warnings and errors. This is a ConsoleHandler
		// so the messages will go to StdOut..
		Logger logger = Logger.getLogger("dr");
		Handler messageHandler = new MessageLogHandler();
		messageHandler.setFilter(new Filter() {
			public boolean isLoggable(LogRecord record) {
				return record.getLevel().intValue() < Level.WARNING.intValue();
			}
		});
		logger.addHandler(messageHandler);

		// Add a handler to handle warnings and errors. This is a ConsoleHandler
		// so the messages will go to StdErr..
		Handler errorHandler = new ConsoleHandler();
		errorHandler.setLevel(Level.WARNING);
		logger.addHandler(errorHandler);

		logger.setUseParentHandlers(false);

		try {
			String fileName = inputFile.getName();

			FileReader fileReader = new FileReader(inputFile);

			XMLParser parser = new BeastParser(new String[] {fileName}, verbose);

			if (consoleApp != null) {
				consoleApp.parser = parser;
			}

			Logger.getLogger("dr.apps.beast").info("Starting cold chain plus hot chains with temperatures: ");
            for (int i = 1; i < chainTemperatures.length; i++) {
                Logger.getLogger("dr.apps.beast").info("Hot Chain " + i + ": " + chainTemperatures[i]);
            }

			Logger.getLogger("dr.apps.beast").info("Parsing XML file: " + fileName);

            ObjectStore store = parser.parse(fileReader, false);
            Iterator iter = store.getIdSet().iterator();
            while (iter.hasNext()) {
                Object id = iter.next();
	            Object obj = store.getObjectById(id);
                if (obj instanceof MCMC) {
                    chains[0] = (MCMC)obj;
                    break;
                }
            }
			if (chains[0] == null) {
				throw new dr.xml.XMLParseException("BEAST XML file is missing an MCMC element");
			}
            fileReader.close();

            chainTemperatures[0] = 1.0;

            for (int i = 1; i < chainCount; i++) {
                fileReader = new FileReader(inputFile);

				messageHandler.setLevel(Level.OFF);
                parser = new BeastParser(new String[] {fileName}, verbose);

			    store = parser.parse(fileReader, false);
	            iter = store.getIdSet().iterator();
                while (iter.hasNext()) {
	                Object id = iter.next();
		            Object obj = store.getObjectById(id);
                    if (obj instanceof MCMC) {
                        chains[i] = (MCMC)obj;
                        break;
                    }
                }
	            if (chains[i] == null) {
		            throw new dr.xml.XMLParseException("BEAST XML file is missing an MCMC element");
	            }
                fileReader.close();
            }
			messageHandler.setLevel(Level.ALL);

		} catch (IOException ioe) {
			System.err.println();
			System.err.println("File error:");
			System.err.println(ioe.getMessage());
		} catch (org.xml.sax.SAXParseException spe) {
			System.err.println();
			if (spe.getMessage() != null && spe.getMessage().equals("Content is not allowed in prolog")) {
				System.err.println("Parsing error - the input file is not a valid XML file.");
			} else {
				System.err.println("Parsing error - poorly formed XML (possibly not an XML file):");
				System.err.println(spe.getMessage());
			}
		} catch (org.w3c.dom.DOMException dome) {
			System.err.println();
			System.err.println("Parsing error - poorly formed XML:");
			System.err.println(dome.getMessage());
		} catch (dr.xml.XMLParseException pxe) {
			if (pxe.getMessage() != null && pxe.getMessage().equals("Unknown root document element, beauti")) {
				System.err.println();
				System.err.println("The file you just tried to run in BEAST is actually a BEAUti document.");
				System.err.println("Although this uses XML, it is not a format that BEAST understands. ");
				System.err.println("These files are used by BEAUti to save and load your settings so that ");
				System.err.println("you can go back and alter them. To generate a BEAST file you must ");
				System.err.println("select the 'Generate BEAST File' option, either from the File menu or ");
				System.err.println("the button at the bottom right of the window.");

			} else {

				System.err.println();
				System.err.println("Parsing error - poorly formed BEAST file:");
				System.err.println(pxe.getMessage());
			}
		} catch (RuntimeException rex) {
			if (rex.getMessage() != null && rex.getMessage().startsWith("The initial model is invalid")) {
				System.err.println();
				System.err.println("The initial model is invalid because state has a zero likelihood.");
				System.err.println("This may be because the initial, random tree is so large that it");
				System.err.println("has an extremely bad likelihood which is being rounded to zero.");
				System.err.println("Alternatively, it may be that the product of starting mutation rate");
				System.err.println("and tree height is extremely small or extremely large. Try to set");
				System.err.println("initial values such that the product is similar to the average");
				System.err.println("pairwise genetic distance between the sequences.");
				System.err.println("For more information go to <http://evolve.zoo.ox.ac.uk/beast/help/>.");

			} else {

				System.err.println();
				System.err.println("Fatal exception (email the authors)");
				rex.printStackTrace(System.out);
			}
		} catch (Exception ex) {
			System.err.println();
			System.err.println("Fatal exception (email the authors)");
			ex.printStackTrace(System.err);
		}

        MCMCMC mc3 = new MCMCMC(chains, options);
        Thread thread = new Thread(mc3);
        thread.start();

	}

	public static void printTitle() {

		System.out.println("+-----------------------------------------------\\");

		String versionString = "BEAST " + version.getVersionString() + " 2002-2004";
		System.out.print("|");
		int n = 47 - versionString.length();
		int n1 = n / 2;
		int n2 = n1 + (n % 2);
		for (int i = 0; i < n1; i++) { System.out.print(" "); }
		System.out.print(versionString);
		for (int i = 0; i < n2; i++) { System.out.print(" "); }
		System.out.println("|\\");

		System.out.println("| Bayesian Evolutionary Analysis Sampling Trees ||");
        System.out.println("|           Metropolis-coupled version          ||");

		String buildString = "BEAST Library: " + version.getBuildString();
		System.out.print("|");
		n = 47 - buildString.length();
		n1 = n / 2;
		n2 = n1 + (n % 2);
		for (int i = 0; i < n1; i++) { System.out.print(" "); }
		System.out.print(buildString);
		for (int i = 0; i < n2; i++) { System.out.print(" "); }
		System.out.println("||");

		System.out.println("|       Alexei Drummond and Andrew Rambaut      ||");
		System.out.println("|              University of Oxford             ||");
		System.out.println("|      http://evolve.zoo.ox.ac.uk/Beast/        ||");
		System.out.println("\\-----------------------------------------------\\|");
		System.out.println(" \\-----------------------------------------------\\");
	}

	public static void printHeader() {
		System.out.println(" +-----------------------------------------------+");
		System.out.println(" | Components created by:                        |");
		System.out.println(" |       Alexei Drummond                         |");
		System.out.println(" |       Roald Forsberg                          |");
		System.out.println(" |       Oliver Pybus                            |");
		System.out.println(" |       Andrew Rambaut                          |");
		System.out.println(" | Thanks to (for use of their code):            |");
		System.out.println(" |       Korbinian Strimmer                      |");
		System.out.println(" |       Oliver Pybus                            |");
		System.out.println(" +-----------------------------------------------+");
		System.out.println();

	}

	public static void printUsage(Arguments arguments) {

		arguments.printUsage("beastmc3", " [-chains <chain-count>] [-delta <temperature-delta>|-temperatures <temperature-list>] [-swap <swap-every>] [-verbose] [-window] [-working] [-help] [<input-file-name>]");
		System.out.println();
		System.out.println("  Example: beastmc3 -chains 3 -delta 1.0 -swap 100 test.xml");
		System.out.println("  Example: beast -temperatures 0.2,0.4 -swap 200 -window test.xml");
		System.out.println();
	}

	//Main method
	public static void main(String[] args) throws IOException {


		Arguments arguments = new Arguments(
			new Arguments.Option[] {
                new Arguments.RealOption("chains", 2, Integer.MAX_VALUE, "number of chains"),
                new Arguments.RealOption("delta", 0.0, Double.MAX_VALUE, "temperature increment parameter"),
                new Arguments.RealArrayOption("temperatures", HOT_CHAIN_COUNT, "a comma-separated list of the hot chain temperatures"),
                new Arguments.IntegerOption("swap", 1, Integer.MAX_VALUE, "frequency at which chains temperatures will be swapped"),
				new Arguments.Option("verbose", "verbose XML parsing messages"),
				new Arguments.Option("window", "provide a console window"),
				new Arguments.Option("working", "change working directory to input file's directory"),
				new Arguments.Option("help", "option to print this message")
			});

		try {
			arguments.parseArguments(args);
		} catch (Arguments.ArgumentException ae) {
			System.out.println();
			System.out.println(ae.getMessage());
			System.out.println();
			printTitle();
			printUsage(arguments);
			System.exit(1);
		}

		if (arguments.hasOption("help")) {
			printTitle();
			printUsage(arguments);
			System.exit(0);
		}

        int chainCount = HOT_CHAIN_COUNT + 1;
        if (arguments.hasOption("chains")) {
            chainCount = arguments.getIntegerOption("chains");
        }

        double delta = 1.0;
        if (arguments.hasOption("delta")) {
            if (arguments.hasOption("temperatures")) {
                System.err.println("Either the -delta or the -temperatures option should be used, not both");
                System.err.println();
                printUsage(arguments);
                System.exit(1);
            }
            delta = arguments.getRealOption("delta");
        }

        double[] chainTemperatures = new double[chainCount];
        chainTemperatures[0] = 1.0;
        if (arguments.hasOption("temperatures")) {
            double[] hotChainTemperatures = arguments.getRealArrayOption("temperatures");
            for (int i = 1; i < chainCount; i++) {
                chainTemperatures[i] = hotChainTemperatures[i - 1];
            }
        } else {
            for (int i = 1; i < chainCount; i++) {
                chainTemperatures[i] = 1.0 / (1.0 + (delta * i));
            }
        }

		int swapChainsEvery = 100;
        if (arguments.hasOption("swap")) {
            swapChainsEvery = arguments.getIntegerOption("swap");
        }

		boolean verbose = arguments.hasOption("verbose");
		boolean window = arguments.hasOption("window");
		boolean working = arguments.hasOption("working");

//		if (System.getProperty("dr.app.beast.main.window", "false").toLowerCase().equals("true")) {
//			window = true;
//		}

		BeastConsoleApp consoleApp = null;

		if (window) {
			System.setProperty("com.apple.macos.useScreenMenuBar","true");
			System.setProperty("apple.laf.useScreenMenuBar","true");
			System.setProperty("apple.awt.showGrowBox","true");

			java.net.URL url = BeastMC3.class.getResource("/images/beast.png");
			javax.swing.Icon icon = null;

			if (url != null) {
	           	icon = new javax.swing.ImageIcon(url);
    		}

			String nameString = "BEAST " + version.getVersionString();
			String aboutString = "Bayesian Evolutionary Analysis Sampling Trees\n" +
									version.getVersionString() + "2002-2004 Alexei Drummond & Andrew Rambaut\n" +
									"University of Oxford";

			consoleApp = new BeastConsoleApp(nameString, aboutString, icon);
		}

		String inputFileName = null;

		String[] args2 = arguments.getLeftoverArguments();

		if (args2.length > 1) {
			System.err.println("Unknown option: " + args2[1]);
			System.err.println();
			printUsage(arguments);
			System.exit(1);
		}

		File inputFile = null;

		if (args2.length > 0) {
			inputFileName = args2[0];
			inputFile = new File(inputFileName);
		}

		if (inputFileName == null) {
			// No input file name was given so throw up a dialog box...
			inputFile = Utils.getLoadFile("BEAST "+version.getVersionString()+" - Select XML input file");
		}

		if (inputFile != null && working) {
			System.setProperty("user.dir", inputFile.getParent());
		}

		printTitle();
		printHeader();
		new BeastMC3(chainTemperatures, swapChainsEvery, inputFile, consoleApp, verbose);
	}
}

